import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
from transformers import SwinForImageClassification
import io

class EnsembleModel(nn.Module):
    def __init__(self, num_classes=49, weight_swin=0.5, weight_efficient=0.5):
        super(EnsembleModel, self).__init__()
        self.swin = SwinForImageClassification.from_pretrained("microsoft/swin-base-patch4-window7-224", ignore_mismatched_sizes=True)
        self.efficientnet = models.efficientnet_b7(weights="IMAGENET1K_V1")
        if isinstance(self.efficientnet.classifier, nn.Sequential):
            for layer in self.efficientnet.classifier:
                if isinstance(layer, nn.Linear):
                    in_features = layer.in_features
                    break
        else:
            in_features = self.efficientnet.classifier.in_features
        self.efficientnet.classifier = nn.Linear(in_features, num_classes)
        self.fc_swin = nn.Linear(1000, num_classes)
        self.weight_swin = weight_swin
        self.weight_efficient = weight_efficient

    def forward(self, x):
        swin_out = self.swin(x).logits
        swin_out = self.fc_swin(swin_out)
        efficient_out = self.efficientnet(x)
        combined_output = self.weight_swin * swin_out + self.weight_efficient * efficient_out
        return combined_output

def load_model(model_path="best_model.pth"):
    model = EnsembleModel().cuda()
    model.load_state_dict(torch.load(model_path))
    model.eval()
    return model

def preprocess_image(image: Image.Image):
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    return transform(image).unsqueeze(0).cuda()

def predict_image(model, image_bytes):
    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    image_tensor = preprocess_image(image)
    with torch.no_grad():
        output = model(image_tensor)
        _, predicted_class = torch.max(output, 1)
    return predicted_class.item()